from datetime import datetime
import tushare as ts
import os


def tushare_client():
    token = os.getenv('TUSHARE_TOKEN')
    if not token:
        print('No env TUSHARE_TOKEN found!')
        exit(1)
    ts.set_token(token)
    pro = ts.pro_api()
    return pro


def current_date():
    return datetime.now().strftime('%Y%m%d')


def date_to_dot_format(str_date):
    """
    yyyymmdd格式的日期转为yyyy.mm.dd格式
    :param str_date: str, yyyymmdd格式的日期
    :return: str, yyyy.mm.dd格式的日期
    """
    assert len(str_date) == len('20201231')
    return str_date[:4] + '.' + str_date[4:6] + '.' + str_date[6:]


def date_to_bar_format(str_date):
    """
    yyyymmdd格式的日期转为yyyy-mm-dd格式
    :param str_date: str, yyyymmdd格式的日期
    :return: str, yyyy-mm-dd格式的日期
    """
    assert len(str_date) == len('20201231')
    return str_date[:4] + '-' + str_date[4:6] + '-' + str_date[6:]


def date_from_bar_format(str_date):
    """
    yyyy-mm-dd格式的日期转为yyyymmdd格式
    :param str_date: str, yyyy-mm-dd格式的日期
    :return: str, yyyymmdd格式的日期
    """
    assert len(str_date) == len('2020-12-31')
    return str_date.replace('-', '')


def date_from_dot_format(str_date):
    """
    yyyy.mm.dd格式的日期转为yyyymmdd格式
    :param str_date: str, yyyy.mm.dd格式的日期
    :return: str, yyyymmdd格式的日期
    """
    assert len(str_date) == len('2020.12.31')
    return str_date.replace('.', '')

def date_to_timestamp(str_date):
    """
    yyyymmdd格式的日期转为时间戳
    :param str_date:
    :return:
    """
    assert len(str_date) == len('20201231')
    # 也可以用 return datetime.strptime(str_date, "%Y%m%d").timestamp()
    return datetime(
        int(str_date[:4]),
        int(str_date[4:6]),
        int(str_date[6:])
    ).timestamp()
